[https://nvbugs/6275959][fix] In runJITWarmupGridIfRequested, std::min the warmup-only maxBatchSize against… - #15221
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughThe attention warmup now skips TRTLLM-Gen FMHA JIT warmup when ChangesFMHA warmup workload guard
Estimated code review effort: 1 (Trivial) | ~5 minutes Merge Risk: ⚪ Minimal · up to The change limits only startup JIT warmup candidate sizes while preserving runtime engine maxima, so no actionable merge-blocking risk remains beyond normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
84f1d06 to
00034fc
Compare
f75b6f1 to
aab26e7
Compare
dad8280 to
2eda7da
Compare
2eda7da to
5533832
Compare
|
The PR description and the actual diff describe two different fixes. Flagging before approval. 1. Description vs. code mismatch The description says the fix is in C++ The actual diff does something different: in Python So the description promises "clamped warmup", but the code implements "skip warmup entirely above a threshold". The description (auto-generated by repair-bot) should be rewritten to match the real change before merge, otherwise future readers/bisects will be misled. 2. Issues with the current implementation
What is fine
Suggestion: the description's "clamped warmup" approach would be preferable — it avoids the timeout while keeping warmup coverage for common shapes. At minimum, please fix the description to reflect the skip-based implementation before merge. |
|
Please have a check on the above issues. If it is intended, I will approve this. |
…nd config Review follow-up on NVIDIA#16614 (yingguo-trt, fredricz-20070104): the consolidation dropped test_disaggregated_logprobs_serving, which is the resident regression guard for NVBug 5926823 (disagg + streaming + logprobs). Its API-level coverage is unique -- streaming vs. non-streaming logprobs consistency across both the completions and chat APIs, plus chat top_logprobs across the ctx/gen boundary. test_disaggregated_single_gpu.py::test_disaggregated_logprobs only covers the core prefill->decode propagation, so deleting this would be a test escape. The test was never UCX-specific; it only reused the UCX yaml. Restore it on a new DEFAULT-backend config and drop the explicit TRTLLM_USE_UCX_KVCACHE pin, which keeps the de-UCX-specialization goal of this PR intact. UCX_TLS is kept, matching every other non-pinned test in this file. The waive for the open bug (NVBugs 6275959, fix in flight in NVIDIA#15221) is restored with it. Restores the numpy import: np.isclose in this test is its only user. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
9d74c00 to
de1f1f8
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 1491-1505: Replace the early return in the FMHA warmup flow with
capped dimensions for the C++ warmup-grid invocation, limiting batch size to 256
and sequence length to 16384 while preserving self.batch_size and
self.max_seq_len for serving. Define named constants for these caps and derive
the candidate-grid workload limit from their product, then pass the capped
values to trtllm_gen_fmha_jit_warmup().
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 250c7d59-7c26-4d53-8405-f5f2968984d8
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
| # The C++ TRTLLM-Gen FMHA JIT warmup enumerates a (batchSize x seqLenKv) cartesian | ||
| # grid sized by engine maxima. PR #15305 densified the candidate lists, so for | ||
| # long-context configs (e.g., max_num_requests=2048, max_seq_len=131072 in | ||
| # disagg_config_ctxtp2_gentp2_llama31_8b_ucx.yaml), the grid produces thousands of | ||
| # NVRTC compilations that exceed the 600s server-start timeout. Skip the warmup | ||
| # whenever the product would blow the budget; any kernel not pre-warmed JIT-compiles | ||
| # lazily on first request, which is correct (just slower for that one request). | ||
| # The threshold matches the pre-PR #15305 effective grid size. | ||
| max_warmup_workload = self.batch_size * self.max_seq_len | ||
| if max_warmup_workload > 256 * 16384: | ||
| logger.info( | ||
| f"Skipping TRTLLM-Gen FMHA JIT warmup: engine config " | ||
| f"(max_batch_size={self.batch_size}, max_seq_len={self.max_seq_len}) " | ||
| f"would produce too many warmup grid points") | ||
| return |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
Replace the all-or-nothing skip with capped FMHA warmup.
This branch returns before trtllm_gen_fmha_jit_warmup() runs. It leaves the TRTLLM-Gen FMHA grid unwarmed and moves NVRTC compilation into live request handling. That conflicts with the PR objective to clamp warmup dimensions to batch size 256 and sequence length 16384 while preserving runtime maxima.
Pass capped dimensions to the C++ warmup-grid path instead of skipping the entire warmup. Keep self.batch_size and self.max_seq_len unchanged for serving. Derive the limit from the actual candidate-grid budget and replace the inline 256 * 16384 values with named constants.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 1491 - 1505,
Replace the early return in the FMHA warmup flow with capped dimensions for the
C++ warmup-grid invocation, limiting batch size to 256 and sequence length to
16384 while preserving self.batch_size and self.max_seq_len for serving. Define
named constants for these caps and derive the candidate-grid workload limit from
their product, then pass the capped values to trtllm_gen_fmha_jit_warmup().
|
/bot run --only-qa-verify test disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] |
|
PR_Github #64171 [ run ] triggered by Bot. Commit: |
|
PR_Github #64171 [ run ] completed with state |
StanleySun639
left a comment
StanleySun639
left a comment
There was a problem hiding this comment.
LGTM if the target test case can pass.
…nd config Review follow-up on NVIDIA#16614 (yingguo-trt, fredricz-20070104): the consolidation dropped test_disaggregated_logprobs_serving, which is the resident regression guard for NVBug 5926823 (disagg + streaming + logprobs). Its API-level coverage is unique -- streaming vs. non-streaming logprobs consistency across both the completions and chat APIs, plus chat top_logprobs across the ctx/gen boundary. test_disaggregated_single_gpu.py::test_disaggregated_logprobs only covers the core prefill->decode propagation, so deleting this would be a test escape. The test was never UCX-specific; it only reused the UCX yaml. Restore it on a new DEFAULT-backend config and drop the explicit TRTLLM_USE_UCX_KVCACHE pin, which keeps the de-UCX-specialization goal of this PR intact. UCX_TLS is kept, matching every other non-pinned test in this file. The waive for the open bug (NVBugs 6275959, fix in flight in NVIDIA#15221) is restored with it. Restores the numpy import: np.isclose in this test is its only user. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
…nd config Review follow-up on NVIDIA#16614 (yingguo-trt, fredricz-20070104): the consolidation dropped test_disaggregated_logprobs_serving, which is the resident regression guard for NVBug 5926823 (disagg + streaming + logprobs). Its API-level coverage is unique -- streaming vs. non-streaming logprobs consistency across both the completions and chat APIs, plus chat top_logprobs across the ctx/gen boundary. test_disaggregated_single_gpu.py::test_disaggregated_logprobs only covers the core prefill->decode propagation, so deleting this would be a test escape. The test was never UCX-specific; it only reused the UCX yaml. Restore it on a new DEFAULT-backend config and drop the explicit TRTLLM_USE_UCX_KVCACHE pin, which keeps the de-UCX-specialization goal of this PR intact. UCX_TLS is kept, matching every other non-pinned test in this file. The waive for the open bug (NVBugs 6275959, fix in flight in NVIDIA#15221) is restored with it. Restores the numpy import: np.isclose in this test is its only user. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
…nd config Review follow-up on NVIDIA#16614 (yingguo-trt, fredricz-20070104): the consolidation dropped test_disaggregated_logprobs_serving, which is the resident regression guard for NVBug 5926823 (disagg + streaming + logprobs). Its API-level coverage is unique -- streaming vs. non-streaming logprobs consistency across both the completions and chat APIs, plus chat top_logprobs across the ctx/gen boundary. test_disaggregated_single_gpu.py::test_disaggregated_logprobs only covers the core prefill->decode propagation, so deleting this would be a test escape. The test was never UCX-specific; it only reused the UCX yaml. Restore it on a new DEFAULT-backend config and drop the explicit TRTLLM_USE_UCX_KVCACHE pin, which keeps the de-UCX-specialization goal of this PR intact. UCX_TLS is kept, matching every other non-pinned test in this file. The waive for the open bug (NVBugs 6275959, fix in flight in NVIDIA#15221) is restored with it. Restores the numpy import: np.isclose in this test is its only user. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
…nd config Review follow-up on NVIDIA#16614 (yingguo-trt, fredricz-20070104): the consolidation dropped test_disaggregated_logprobs_serving, which is the resident regression guard for NVBug 5926823 (disagg + streaming + logprobs). Its API-level coverage is unique -- streaming vs. non-streaming logprobs consistency across both the completions and chat APIs, plus chat top_logprobs across the ctx/gen boundary. test_disaggregated_single_gpu.py::test_disaggregated_logprobs only covers the core prefill->decode propagation, so deleting this would be a test escape. The test was never UCX-specific; it only reused the UCX yaml. Restore it on a new DEFAULT-backend config and drop the explicit TRTLLM_USE_UCX_KVCACHE pin, which keeps the de-UCX-specialization goal of this PR intact. UCX_TLS is kept, matching every other non-pinned test in this file. The waive for the open bug (NVBugs 6275959, fix in flight in NVIDIA#15221) is restored with it. Restores the numpy import: np.isclose in this test is its only user. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
de1f1f8 to
5d3ccfe
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
5d3ccfe to
792d717
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
ae966dd to
eb82de8
Compare
…ngine configs PR NVIDIA#14851 added a TRTLLM-Gen FMHA JIT warmup that enumerates a cartesian grid of (batchSize, seqLenQ, seqLenKv) sized by the engine maxima, eagerly compiling NVRTC kernels for each combo. PR NVIDIA#15305 then densified the candidate lists to catch missing kernels. For long-context disagg configs such as disagg_config_ctxtp2_gentp2_llama31_8b_ucx.yaml (max_num_requests=2048, max_seq_len=131072), the densified grid contains thousands of points and the NVRTC compilation time exceeds the 600s wait_for_disagg_server_ready timeout in test_disaggregated_logprobs_serving[llama-3.1-8b-instruct]. Skip the warmup at the Python entry point when the engine maxima product would produce a problematic grid. The runtime kernel selection path is unchanged -- any kernel that would have been warmed up will JIT-compile lazily on first request instead. This restores the documented pre-PR NVIDIA#14851 behavior for oversized configs without affecting the warmup benefit for regular configs. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
eb82de8 to
1462190
Compare
9 participants
Summary
Test plan
Links
Dev Engineer Review
_run_attention_warmupskips TRTLLM-Gen FMHA JIT warmup whenbatch_size * max_seq_len > 256 * 16384.QA Engineer Review
tests/integration/test_lists/waives.txtre-enables the disaggregated log-probability serving test forllama-3.1-8b-instruct.